# -*- coding:utf8 -*-
"""
先对新闻进行LDA预分类后，再进行单独类别的处理和标注
"""
from __future__ import print_function, unicode_literals
import os
from gensim.corpora import Dictionary
from gensim.models.ldamulticore import LdaMulticore
from ossys_utility.dirfunc import mkdir_if_not_exist
from coverkg.config import MODEL_DIR


LDA_CLASSIFY_PPATH = os.path.join(MODEL_DIR, 'LDA_CLASSIFY')
mkdir_if_not_exist(LDA_CLASSIFY_PPATH)


class LdaClassify(object):
    """
    LDA分类
    """
    def __init__(self, tag=None):
        """
        初始化
        :param tag:
        """
        self.tag = tag
        self.classify = None
        self.dictionary = None

    def _default_path(self):
        """
        默认保存位置
        :return:
        """
        if self.tag is None:
            tag = 'DEFAULT'
        else:
            tag = self.tag

        save_dir = os.path.join(LDA_CLASSIFY_PPATH, tag)
        mkdir_if_not_exist(save_dir)

        return os.path.join(save_dir, 'GenSim.LdaModel.model')

    def load(self, path_base=None):
        """
        加载模型
        :param path_base:
        :return:
        """
        if path_base is None:
            path_base = self._default_path()
        path_lda = path_base + '.lda'
        path_dict = path_base + '.dict'

        self.classify = LdaMulticore.load(path_lda)
        self.dictionary = Dictionary.load(path_dict)

    def save(self, path_base=None):
        """
        保存模型
        :param path_base:
        :return:
        """
        if path_base is None:
            path_base = self._default_path()
        path_lda = path_base + '.lda'
        path_dict = path_base + '.dict'

        self.classify.save(path_lda)
        self.dictionary.save(path_dict)

    def predict(self, content):
        """
        获取分类信息
        :param content:
        :return:
        """
        vec = self.dictionary.doc2bow(content)
        return self.classify[vec]

    def train(self, texts, n_topic):
        """
        训练
        :param texts:
        :param n_topic:
        :return:
        """
        dictionary = Dictionary(texts)
        corpus = [dictionary.doc2bow(text) for text in texts]
        self.dictionary = dictionary
        self.classify = LdaMulticore(corpus, id2word=dictionary, num_topics=n_topic)


if __name__ == '__main__':
    def lda_classify_test():
        news = [
            ['日你', '妈妈'],
            ['阿妈', '爱', '你'],
            ['你', '给', '我', '滚'],
            ['太阳'],
            ['月亮'],
            ['风险', '投资'],
            ['投资', '运营'],
            ['黄', '赌', '毒'],
            ['于娜娜'],
            ['太极']
        ]
        lda_classify = LdaClassify('test')
        lda_classify.train(news, 3)
        print(lda_classify.classify.print_topics(20))
        print(lda_classify.predict(['风险', '与', '南安']))
        lda_classify.save()

    lda_classify_test()